# Copyright (c) 2019 Presto Labs Pte. Ltd.
# Author: leon

import datetime
import json
import logging
from typing import List, Union

from tornado import gen

import coin.flow.topic_map as topic_map
from coin.base.param_util import to_list
from coin.exchange.base.ws.subscriber import WebsocketJsonSubscriberBase
from coin.exchange.bybit.kr_rest.futures_product import BybitFuturesProduct


class BybitSubscriber(WebsocketJsonSubscriberBase):
  def __init__(self,
               queue,
               group_id: str,
               worker_id: str,
               symbol: Union[List[str], str],
               close_callback=None,
               logger=None):
    super().__init__(close_callback, logger or logging.getLogger(__name__))
    self._symbol = to_list(symbol)

    self._register_url('wss://stream.bybit.com/realtime')
    self._publisher = None
    self._init_publisher(queue, group_id, worker_id)
    self._last_pong_time = None

  def _init_publisher(self, queue, group_id: str, worker_id: str):
    topic_depth = topic_map.search_topic_id('bybit_ws_depth_%s%s' % (group_id, worker_id))[0]
    topic_trade = topic_map.search_topic_id('bybit_ws_trade_%s%s' % (group_id, worker_id))[0]
    topic_instrument = topic_map.search_topic_id('bybit_ws_instrument_%s%s' %
                                                 (group_id, worker_id))[0]

    self._publisher = {
        'depth': queue.get_writer(topic_id=topic_depth[1]),
        'trade': queue.get_writer(topic_id=topic_trade[1]),
        'instrument': queue.get_writer(topic_id=topic_instrument[1]),
    }

    self._enable_topic_event_publisher(queue,
                                       'bybit_ws_topic_event',
                                       [topic_depth[0], topic_trade[0], topic_instrument[0]])
    self._add_teardown_callback(self._teardown_publisher)

  @gen.coroutine
  def on_open(self):
    self._ioloop.add_callback(self._send_subscribe_request)

  @gen.coroutine
  def _send_subscribe_request(self):
    products = [BybitFuturesProduct.FromStrNativeProduct(symbol) for symbol in self._symbol]
    args = ['orderBookL2_25.%s' % product.native_subscription_symbol for product in products]
    subscribe_depth_str = json.dumps({"op": "subscribe", "args": args})
    yield self._ws.write_message(subscribe_depth_str)
    yield gen.sleep(0.5)

    args = ['trade']
    subscribe_trade_str = json.dumps({"op": "subscribe", "args": args})
    yield self._ws.write_message(subscribe_trade_str)
    yield gen.sleep(0.5)

    args = ['instrument_info.100ms.%s' % product.native_subscription_symbol for product in products]
    subscribe_instrument_str = json.dumps({"op": "subscribe", "args": args})
    yield self._ws.write_message(subscribe_instrument_str)
    yield gen.sleep(0.5)

    self._publish_topic_event('START')
    self._logger.info('Subscribe done')

    self._last_pong_time = datetime.datetime.now()
    while self._state == self.State.RUNNING:
      if (datetime.datetime.now() - self._last_pong_time >= datetime.timedelta(minutes=2)):
        self._logger.error('No pong has been received since %s. Closing...'
                           % str(self._last_pong_time))
        self._notify_close()
        break

      self.ping()
      yield gen.sleep(30)

  def _teardown_publisher(self):
    if self._publisher:
      for _, publisher in self._publisher.items():
        publisher.close()
      self._publisher = None

  def on_message(self, timestamp: int, message: bytes):
    if self.stopped:
      return

    if 'ret_msg' in message and message['ret_msg'] == 'pong':
      self._last_pong_time = datetime.datetime.now()
      self._logger.debug('pong')
    elif 'success' in message:
      self._publish_topic_event(message)
    elif 'topic' in message:
      topic = message['topic']
      if topic.startswith('orderBookL2_25'):
        topic = 'depth'
      elif topic.startswith('trade'):
        topic = 'trade'
      elif topic.startswith('instrument_info.100ms'):
        topic = 'instrument'
      else:
        self._logger.error('Unknown topic: %s' % message)
        return
      self._publisher[topic].write(timestamp, message)
    else:
      self._logger.error('Unknown type: %s' % message)

  @gen.coroutine
  def ping(self):
    yield self._ws.write_message(json.dumps({"op": "ping"}))
    self._logger.debug('ping')
